import numpy as np
import pandas as pd


def shannon_entropy():
    """
    香农熵
    :return:
    """
    p = 1 / 2
    return -((p * np.log2(p)) + (p * np.log2(p)))


def calEnt(dataSet):
    """
    计算数据集的香农熵
    :param dataSet:
    :return:
    """
    n = dataSet.shape[0]  # 数据集总行数
    # [[不是, 3], [是, 2]]
    iset = dataSet.iloc[:, -1].value_counts()  # 标签的所有类别
    p = iset / n  # 每一类标签所占比
    ent = (-p * np.log2(p)).sum()  # 计算信息熵
    return ent


row_data = {'是否陪伴': [0, 0, 0, 1, 1],
            '是否玩游戏': [1, 1, 0, 1, 1],
            '渣男': ['是', '是', '不是', '不是', '不是']}
dataSet = pd.DataFrame(row_data)

# 计算全体数据的信息熵——根据标签列去进行计算
# 0.9709505944546686
print(calEnt(dataSet))


# 选择最优的列进行切分
def bestSplit(dataSet):
    baseEnt = calEnt(dataSet)  # 计算原始熵
    bestGain = 0  # 初始化信息增益
    axis = -1  # 初始化最佳切分列，标签列
    for i in range(dataSet.shape[1] - 1):  # 对特征的每一列进行循环
        levels = dataSet.iloc[:, i].value_counts().index  # 提取出当前列的所有取值
        ents = 0  # 初始化子节点的信息熵
        for j in levels:  # 对当前列的每一个取值进行循环
            childSet = dataSet[dataSet.iloc[:, i] == j]  # 某一个子节点的dataframe
            ent = calEnt(childSet)  # 计算某一个子节点的信息熵
            ents += (childSet.shape[0] / dataSet.shape[0]) * ent  # 计算当前列的信息熵
        print('第{}列的信息熵为{}'.format(i, ents))
        infoGain = baseEnt - ents  # 计算当前列的信息增益
        print('第{}列的信息增益为{}\n'.format(i, infoGain))
        if (infoGain > bestGain):
            bestGain = infoGain  # 选择最大信息增益
            axis = i  # 最大信息增益所在列的索引
    print("第{}列为最优切分列".format(axis))
    return axis

print(bestSplit(dataSet)) # 0

def mySplit(dataSet, axis, value):
  """
  函数功能：按照给定的列划分数据集
  参数说明：
  dataSet：原始数据集
  axis：指定的列索引
  value：指定的属性值
  return：redataSet：按照指定列索引和属性值切分后的数据集
  """
  col = dataSet.columns[axis]
  redataSet = dataSet.loc[dataSet[col] == value, :].drop(col, axis=1)
  return redataSet

# 验证函数：以axis=0，value=1为例
print(mySplit(dataSet, 0, 1))

